# -*- coding: utf-8  -*-
import sys, os, re

import configparser
import pandas as pd 
import numpy as np
import nibabel as nib
import copy


cf = configparser.ConfigParser(allow_no_value=False, comment_prefixes=('#', ';'))
cf.read("config.ini", encoding='utf-8')  # 读取配置文件，如果写文件的绝对路径，就可以不用os模块


mycf = configparser.ConfigParser(allow_no_value=False, comment_prefixes=('#', ';'))
mycf.read("dataset.ini", encoding='utf-8')


# 保存文件信息
"""
{
    "别名": {
        {"path": 文件路径, "data": dataframe数据, "cols": 列名}
    },
}
"""
files = {}

# 保存合并后的文件图像列名
imagesCols = []

# 剩下的列名
columnsNew = []

# print(cf.options("INFO"))
# print(cf.items())
# print(cf.sections())
# print(cf.get("INFO", "NAme"), cf["INFO"]["NAme"])

def configToDict(config_name, config="INFO"):

    res_dict = {}
    cf = configparser.ConfigParser()
    cf.read(config_name, encoding='utf-8') 
    print(cf.items(config))
    for item in cf.items(config):
        print(item[0], item[1])
        res_dict[item[0]] = item[1]

    return res_dict

# 获取配置文件cf对应section章节的key值
def getConfig(cf, section, key):

    if section not in cf.sections():
        print("配置%s不存在" % section)
        return None

    # keys = [item[0] for item in cf.options[section]]
    if key not in cf.options(section):

        print("配置项%s不存在" % key)

        return None
    
    return cf[section][key]

# 获取合并后的数据列名，例如 gene|RID 得到RID
def getColName(colStr):

    if '|' in colStr:
        myCol = argSplit(colStr, '|')
        # print(files.keys())
        if myCol[0] not in files.keys():
            print("没有%s数据集" % myCol[0])
            return None
        if myCol[1] not in files[myCol[0]]["cols"].keys():
            print("没有%s列名" % myCol[1])
            return None
        return files[myCol[0]]["cols"][myCol[1]]
    elif colStr:
        return colStr

    else:
        return None


# 获取dataName的列名，例如 ["blood|ID", "gene|RID", "gene|date"] 得到[RID,date]
def getColByData(myCOls, dataName):

    selected = []
    for C in myCOls:
        myCol = argSplit(C, '|')
        if myCol[0] == dataName:
            selected.append(myCol[1])

    return selected


# 给重复的列名增加后缀
def addSuffix(fdict, s1, s2):

    list1 = fdict[s1]["data"].columns
    list2 = fdict[s2]["data"].columns
    cols1 = fdict[s1]["cols"]
    cols2 = fdict[s2]["cols"]
    for i in range(len(list1)):
        for j in range(len(list2)):
            if list1[i] == list2[j]:
                cols1[list1[i]] += "_%s" % s1
                cols2[list2[j]] += "_%s" % s2
                print("数据集 <%s> 的 <%s> 列和数据集 <%s> 的 <%s> 列相同，已分别修改为：%s, %s" \
                     % (s1, list1[i], s2, list2[j], cols1[list1[i]], cols2[list2[j]]))
                
                # 修改列名
                fdict[s1]["data"] = files[s1]["data"].rename(columns=cols1)
                fdict[s2]["data"] = files[s2]["data"].rename(columns=cols2)
                

""" 解析函数的参数值 """
def getArg(expStr, arg):

    p = re.compile(r'[(](.*?)[)]', re.S) # 最小匹配 

    # 获取括号内的参数
    args_str = re.findall(p, expStr)

    if len(args_str) > 0:
        argstr = args_str[0]
    else:
        print("参数错误")
        return ""
    
    if argstr == "" or arg not in argstr:
        return ""

    args = argSplit(argstr,',')

    # 提取参数为arg的字符串
    myarg = ""
    for i in args:
        if arg in i and '=' in i:
            item = argSplit(i, '=')
            if arg in item[0]:
                # 返回=号后面的字符串
                myarg = item[1]
                break

    return myarg

# 按照字符分割字符串，并逐一去空格后返回
def argSplit(argStr, sep=','):

    if not isinstance(argStr, str):
        print("参数不为字符串")
        return None

    if not argStr:
        return None
    arr = argStr.split(sep)

    # 去空格
    for t in range(len(arr)):
        arr[t] = arr[t].strip()

    # 去空字符串
    if "" in arr:
        arr.remove("")

    return arr

# 给数据添加图像数据
def addImage(client):

    global imagesCols
    for c in range(len(client)):
        # 当前数据的长度
        curLen = len(client[c]["data"])
        # 遍历表示图像的列名
        for i in imagesCols:
            imageIndex = columnsNew.index(i)
            imagePaths = client[c]["data"][imageIndex]
            imagePaths = argSplit(imagePaths)
            if imagePaths == None:
                continue
            # 处理多张图片

            for imagePath in imagePaths:
                if os.path.exists(imagePath):
                    imgData = nib.load(imagePath)
                    imgData = np.asanyarray(imgData.dataobj)
                    shape = list(imgData.shape)
                    imgData = imgData.reshape(-1).tolist()
                    client[c]["data"].append(imgData)
                    client[c]["headers"]["image"].append({"range": [curLen, curLen+len(imgData)-1], "shape": shape})
                    curLen += len(imgData)

    return client

# 微调数据
def _adjust(data, addStr, df):

    adds = argSplit(addStr)

    for a in adds:
        item = argSplit(a, "@")
        index = int(item[0]) - 1
        cond = item[1]
        rate = float(item[2])

        add_data = df_filter(df, cond).values
        # 按照比例随机抽取数据
        # print(add_data)
        np.random.shuffle(add_data)
        add_data = add_data[:int(rate * add_data.shape[0])]
        # 合并数组
        data[index] = np.vstack((data[index], add_data))

    return data

# 按照条件筛选数据分给客户端
# condition是字符串，是该客户端满足的条件，例如"gene|RID>500 & blood|SITE=10"
def df_filter(df, condition):

    conds = argSplit(condition, '&')
    # print()
    tempdf = df #.astype("str")
    # 按照条件筛选
    for con in conds:
        if '=' in con:
            eq = argSplit(con, '=')
            # print("eq ", eq)
            # 列名转换
            eq[0] = getColName(eq[0])

            # 整型
            if "int" in str(tempdf[eq[0]].dtype):
                eq[1] = int(eq[1])
            # 浮点型
            elif "float" in str(tempdf[eq[0]].dtype):
                eq[1] = float(eq[1])
            tempdf = tempdf[tempdf[eq[0]] == eq[1]]
            # print("eq ", eq, tempdf.shape)
        
        elif '<' in con:
            lt = argSplit(con, '<')
            # 列名转换
            lt[0] = getColName(lt[0])
            # 整型
            if "int" in str(tempdf[lt[0]].dtype):
                lt[1] = int(lt[1])
            # 浮点型
            elif "float" in str(tempdf[lt[0]].dtype):
                lt[1] = float(lt[1])

            tempdf = tempdf[tempdf[lt[0]] < lt[1]]
            # print("< ", lt, tempdf.shape)

        elif '>' in con:
            gt = argSplit(con, '>')
            # print(gt)
            # 列名转换
            gt[0] = getColName(gt[0])

            # 整型
            if "int" in str(tempdf[gt[0]].dtype):
                gt[1] = int(gt[1])
            # 浮点型
            elif "float" in str(tempdf[gt[0]].dtype):
                gt[1] = float(gt[1])
            
            tempdf = tempdf[tempdf[gt[0]] > gt[1]]
            # print("> ", gt, tempdf.shape)
        else:
            print("条件表达式错误")

    return tempdf


# 按照条件筛选数据
# condition是一个列表，是该文件的数据需要满足的条件, 例如[RID>500, SITE<500]
def df_filter_rows(df, condition):

    # 按照条件筛选
    
    tempdf = df #.astype("str")

    for con in condition:

        if '!=' in con:
            neq = argSplit(con, '!=')
            print("neq ", neq)
            # 列名转换
            neq[0] = getColName(neq[0])

            if neq[1] == "null":
                tempdf = tempdf[(tempdf[neq[0]].notnull()) & (tempdf[neq[0]] != "")]
            else:
                # 整型
                if "int" in str(tempdf[neq[0]].dtype):
                    neq[1] = int(neq[1])
                # 浮点型
                elif "float" in str(tempdf[neq[0]].dtype):
                    neq[1] = float(neq[1])
                tempdf = tempdf[tempdf[neq[0]] != neq[1]]
        
        elif '=' in con:
            eq = argSplit(con, '=')
            # print("eq ", eq)
            # 列名转换
            eq[0] = getColName(eq[0])
            # 整型
            if "int" in str(tempdf[eq[0]].dtype):
                eq[1] = int(eq[1])
            # 浮点型
            elif "float" in str(tempdf[eq[0]].dtype):
                eq[1] = float(eq[1])
            tempdf = tempdf[tempdf[eq[0]] == eq[1]]
            # print("eq ", eq, tempdf.shape)
        
        elif '<' in con:

            lt = argSplit(con, '<')
            # 列名转换
            lt[0] = getColName(lt[0])

            # 整型
            if "int" in str(tempdf[lt[0]].dtype):
                lt[1] = int(lt[1])
            # 浮点型
            elif "float" in str(tempdf[lt[0]].dtype):
                lt[1] = float(lt[1])

            tempdf = tempdf[tempdf[lt[0]] < lt[1]]
            print("< ", lt, tempdf.shape)

        elif '>' in con:
            gt = argSplit(con, '>')
            # print(gt)
            # 列名转换
            gt[0] = getColName(gt[0])

            # 整型
            if "int" in str(tempdf[gt[0]].dtype):
                gt[1] = int(gt[1])
            # 浮点型
            elif "float" in str(tempdf[gt[0]].dtype):
                gt[1] = float(gt[1])
            
            tempdf = tempdf[tempdf[gt[0]] > gt[1]]
            # print("> ", gt, tempdf.shape)

        
        else:
            print("条件表达式错误")

    return tempdf




def dfs(df, C, group, selected, n, res):

    if n == len(C):
        
        data = df[df[group[0]] == selected[group[0]]]
        if data.shape[0] == 0:
            return
        # print(selected)
        for s in group[1:]:
            # print(data)
            if data.shape[0] == 0:
                break
            data = data[data[s] == selected[s]]
            
            # data = data.values

        if (data.values).shape[0] > 0:
            res.append(data)

    else:
        for i in range(len(C[n])):
            selected[group[n]] = C[n][i]
            dfs(df, C, group, selected, n+1, res)

def _class(df, group, workers=None, ):
    """
    按照指定列的组合类划分
    """
    res = []

    # print("列名<%s>的种数是%d,需要分配的份数：%s" % (column, len(df[column].value_counts()), workers))

    # C = df[group].value_counts()

    if not group:
        print("group为空")
        return []

    N = len(group)

    selected = {}
    C = []
    for i in group:
        selected[i] = ""
        C.append(list(df[i].value_counts().index))
    
    # print(1111111111111)
    # print(len(df['SITE'].value_counts()))
    # print(len(df['APOE4'].value_counts()))
    dfs(df, C, group, selected, 0, res)

    return res


def _scope(df, workers, rate=None, sort_by=None, ascending=False):
    """
    按照列名的种类划分
    """
    res = []

    if sort_by != None:
         df.sort_values(by=sort_by , ascending=ascending, inplace=True)

    # 按照比例划分
    if rate != None:
        L = df.shape[0]
        cur = 0

        for r in rate:
            
            if r == rate[-1]:
                res.append(df.iloc[cur:])
            else:
                # print("r=%f,l=%d" % (r, L))
                l = int(r*L)
                res.append(df.iloc[cur:cur+l])
                cur += l
    # 平均分
    else:

        n = int(df.shape[0] / workers)

        for i in range(0, df.shape[0], n):

            # 剩下的全部分配
            if i + 2*n > df.shape[0]:
                res.append(df.iloc[i:])
                break

            # 分配第i段数据
            else:
                res.append(df.iloc[i:i+n])
    
    return res



def _constraint(df, workers, conditions):
    """
    按照限制条件划分
    """
    res = []

    # print(conditions)

    n = 0
    for c in conditions:

        tempdf = df_filter(df, c)
        
        res.append(tempdf)
        n += 1
        if n == workers:
            break

    return res


def allocation(block="block1"):
    """ 
    配置场景块为block的场景，返回block配置中对应的数据，类型为list，list中每个元素是Numpy array数组
    """

    print("------ 场景配置开始 ------")
    print("场景名称：%s " % getConfig(mycf, "scenarion", "name"))
    print("场景描述：%s " % getConfig(mycf, "scenarion", "description"))
    print("场景版本：%s " % getConfig(mycf, "scenarion", "version"))
    
    print("读取的场景块：%s " % block)


    global imagesCols
    global columnsNew
    global files

    ######### 读取配置信息 #########

    exp = argSplit(getConfig(mycf, block, "assignment"), '(')[0]

    print("数据划分函数：%s " % exp)

    workers = getConfig(mycf, block, "partition_number")

    filesStr = getConfig(mycf, block, "files")
    # filesStr = mycf[block]["Files"]

    adds = getConfig(mycf, block, "add")

    joinorder = getConfig(mycf, block, "joinorder")

    connection = getConfig(mycf, block, "connection")
    # connection = mycf[block]["connection"]

    columns = getConfig(mycf, block, "columns")

    rows = getConfig(mycf, block, "rows")

    label = getConfig(mycf, block, "label")

    images = getConfig(mycf, block, "image")

    audios = getConfig(mycf, block, "audio")

    videos = getConfig(mycf, block, "video")


    # 分割列名
    columns = argSplit(columns)

    # 分割行筛选条件
    rows = argSplit(rows)

    # 分割多标签
    label = argSplit(label)



    ######### 读取配置信息 #########


    # 合并后的文件
    df = None

    # 保存文件路径以及数据
    fList = argSplit(filesStr, ',')
    for f in fList:
        myfile = argSplit(f, '@')
        mypd = None
        # 如果有部分选择性读取列名
        if columns:
            # 获取该文件选择的列
            selectedCols = getColByData(columns, myfile[0])

            # 如果该文件没有选择任何列，那么全部选择文件的列
            if selectedCols == []:
                mypd = pd.read_csv(myfile[1], low_memory=False, index_col=False)
                # 补充列名 
                columns += [myfile[0] + '|' + c for c in list(mypd.columns)]
            # 只读取选择的列数据
            else:
                mypd = pd.read_csv(myfile[1], usecols=selectedCols, low_memory=False, index_col=False)
        else:
            mypd = pd.read_csv(myfile[1], low_memory=False, index_col=False)
        cols = {}
        
        # "cols": {"源列名"： "新列名"}
        for c in mypd.columns:
            cols[c] = c
        
        # 筛选符合自己的条件
        # 之所以放在这里筛选，而不是合并（merge）之后再筛选，考虑到数据量大，减少不必要的运行时间
        if rows:
            myFilter = []
            for r in rows:
                # print(r)
                
                if myfile[0] in argSplit(r, '|')[0]:
                     # 不需要数据别名，因为这个阶段识别不出数据别名对应的列
                     myFilter.append(argSplit(r, '|')[1])
            if len(myFilter) > 0:
                # 按条件行过滤

                # print(myFilter)
                # print("====过滤前====")
                # print(mypd)
                mypd = df_filter_rows(mypd, myFilter)
                # print("====过滤后====")
                # print(mypd)

                
        files[myfile[0]] = {"path": myfile[1], "data": mypd, "cols": cols}

    # 去除重复列
    key_names = list(files.keys())
    for i in range(len(key_names)):
        for j in range(i+1, len(key_names)):

            addSuffix(files, key_names[i], key_names[j])

    # # 修改列名
    # for k in key_names:
    #     files[k]["data"] = files[k]["data"].rename(columns=files[k]["cols"])

    # print(files[key_names[0]]["cols"])

    print("已配置的文件集:" , key_names)
    # 文件别名列表
    names = argSplit(joinorder, ',')
    # 链接数组, 每个元素都是一个数组，每个数组是2个（左表和右表）若干需要链接的列名
    keys = []
    k = None 
    # 给各个文件数据建立连接
    if connection:
        myConnections = argSplit(connection, ',')
        if len(names) - 1 != len(myConnections):
            print("joinorder和connection数量上不一致，请按照要求填写")
            sys.exit()

        for r in myConnections:
            myKetLists = [[], []]
            colsItem = argSplit(r, '+')
            for item in colsItem:
                myKeys = argSplit(item, '&')
                if len(myKeys) != 2:
                    print("connection格式不正确, 请在配置文件中按照格式填写")
                    sys.exit()
                myKetLists[0].append(getColName(myKeys[0]))
                myKetLists[1].append(getColName(myKeys[1]))

            keys.append(myKetLists)

            # rela = argSplit(r, '&')
            # names.append(rela[0].split('|')[0])
            # for c in rela:
            #     keys.append(getColName(c))

        df = files[names[0]]["data"]

        for i in range(len(names[1:])):
            df_r = files[names[i+1]]["data"]
            k_l = keys[i][0]
            k_r = keys[i][1]
            df = pd.merge(df, df_r, left_on=k_l, right_on=k_r, how='outer')
            
        # 去除重复行
        df.drop_duplicates()

        # print(df.columns)
        # sys.exit()

    # 没有文件合并
    else:
        df = files[files.keys()[0]]["data"]
    

    # 如果选择全部的列，即不填写columns，需要全部添加
    if columns == None:
        columns = []
        for f in files.keys():
            colsTem = copy.deepcopy(list(files[f]["data"].columns))
            print((colsTem))
            for i in range(len(colsTem)):
                colsTem[i] = f + colsTem[i]
            columns.append(colsTem)

    print('选择的数据列是：')
    print(columns)

    # print(df.columns)
    tem = []
    for i in range(len(columns)):
        if getColName(columns[i]):
            tem.append(getColName(columns[i]))

    # 合并之后新表中已选择的列名
    columnsNew = tem
    # if columns:
    #     print("已选择的列：", columns)
    #     df = df[columns]

    print("\n合并后数据的维度:", df.shape)

    # 返回的数据
    res = []

    # 联邦学习场景，返回多个客户端数据
    
    if workers and int(workers) > 1:

        # 按照列的种类划分
        if exp == "class":

            expStr = getConfig(mycf, block, "assignment")
            # expStr = mycf[block]["Expression"]

            group = getArg(expStr, "group")
            group = argSplit(group, '+')
            myGroup = []
            for g in group:
                cn = getColName(g)

                if cn not in columnsNew:
                    print("class函数中参数的列<%s>不在数据集选择的列当中" % g)
                    return []

                if cn:
                    myGroup.append(cn)

            # workers = int(workers) if workers != "" else len(df[group].value_counts())
            
            print("mygpoup", myGroup)
            res = _class(df, group=myGroup, workers=workers)

        # 分段划分
        elif exp == "scope":

            expStr = getConfig(mycf, block, "assignment")
            

            if workers == "":

                print("使用scope函数请指定partition_number")

                return []    
                
            workers = int(workers)

            rate = getConfig(mycf, block, "rate")
            # rate = mycf[block]["rate"]

            if rate:
                rate = argSplit(rate, ',')

            ascending = getArg(expStr, "order")

            sort_by = getArg(expStr, "sort_by")
            

            sort_by = getColName(sort_by)

            print("调用排序类型：%s" % ascending)

            if ascending == 'asc':
                ascending = True
            elif ascending == 'desc': 
                ascending = False
            else:
                print("order 参数错误")
                ascending = False

            if rate:
                if len(rate) != workers:
                    print("划分份数和比例数不一致！")
                    return []
                for i in range(len(rate)):
                    rate[i] = float(rate[i])
                
                print("rate=", rate)
                res = _scope(df, workers=workers, rate=rate, sort_by=sort_by, ascending=ascending)
            else:
                res = _scope(df, workers=workers, sort_by=sort_by, ascending=ascending)

        # 按照限制条件划分
        elif exp == "constraint":

            expStr = getConfig(mycf, block, "assignment")
            

            if workers == "":

                print("使用constraint函数请指定partition_number")

                return []    
                
            workers = int(workers)

            rate = getConfig(mycf, block, "rate")

            if rate:
                rate = argSplit(rate, ',')
            
            conditions = getConfig(mycf, block, "conditions")

            if conditions:

                conditions = argSplit(conditions, ',')

                if len(conditions) != workers:
                    print("限制条件数和比例数不一致！")
                    return []
                print("condition=", conditions)
                res = _constraint(df, workers=workers, conditions=conditions)
            else:
                
                print("请填写限制条件")


        # 微调数据
        if adds != None:
            print("启用数据微调...>>>")
            print(adds)
            print("完成数据微调...<<<")
            _adjust(res, adds, df)

    # 单机模式
    else:

        res = [df]



    # 提取label
    myLabel = []
    labelList = []
    if label:
        myLabel = [getColName(c) for c in label]        
        for l in myLabel:
            if l in columnsNew:
                columnsNew.remove(l)
            else:
                print("****** 标签<%s>没有包含在选择的数据列中")
                myLabel.remove(l)

    if len(myLabel) > 0:
        for i in range(len(res)):
            clientLabel = res[i][myLabel].values.tolist()
            labelList.append(clientLabel)
            res[i] = res[i][columnsNew]

    result = []

    # 数据格式转换成字字典
    for c in range(len(res)):
        res[c] = res[c].values
        myClient = []
        for i in range(len(res[c])):
            lableData = []
            # 有label数据
            if len(labelList) > 0:
                lableData = labelList[c][i]
            simple = {
                "data": res[c][i].tolist(), 
                "label": lableData,
                "labelName": myLabel,
                "headers": {
                    "basicData":{"range": [0, len(res[c][i].tolist()) - 1]},
                    "image": [],
                    "audio": [],
                    "video": []   
                }
            }
            myClient.append(copy.deepcopy(simple))
        
        result.append(myClient)

    del res 


    ######### 读取文件信息 #########

    # 有图片信息
    print('iamge field: ', (images))
 
    # print((res[0][0]))
    # 图片
    if images:
        imagesList = argSplit(images) 
        imagesCols = [getColName(img) for img in imagesList]
        for i in range(len(imagesCols)):
            if imagesCols[i] not in columnsNew:
                
                print("配置文件没选择 %s 列名" % imagesCols[i])
                del imagesCols[i]


        # # 对每个客户端数据
        # for client in result:
        #     for c in range(len(client)):
        #         # 当前数据的长度

        #         curLen = len(client[c]["data"])
        #         # 遍历表示图像的列名
        #         for i in imagesCols:
        #             imageIndex = columnsNew.index(i)
        #             imagePaths = client[c]["data"][imageIndex]
        #             imagePaths = argSplit(imagePaths)
        #             if imagePaths == None:
        #                 continue
        #             # 处理多张图片
        #             for imagePath in imagePaths:
        #                 if os.path.exists(imagePath):
        #                     imgData = nib.load(imagePath)
        #                     imgData = np.asanyarray(imgData.dataobj)
        #                     shape = list(imgData.shape)
        #                     imgData = imgData.reshape(-1).tolist()
        #                     client[c]["data"].append(imgData)
        #                     client[c]["headers"]["image"].append({"range": [curLen, curLen+len(imgData)-1], "shape": shape})
        #                     curLen += len(imgData)



    ######### 读取文件信息 #########


    return result







if __name__ == '__main__':

    # columns = getConfig(mycf, "block1", "columns")


    B = 3
    mydata = allocation()

    print(mydata[0][:3])

    # for i in range(len(mydata)):
    #     print("第%d个客户端的数据量为：%d" % (i+1, len(mydata[i])))

    #     for k in range(0, len(mydata[i]), B):
    #         batData = addImage(mydata[i][k:k+B])
    #         print("1st headers infomation: ")
    #         print(batData[0]["headers"])
    #         # print(batData[1]["headers"])
    #         # print(batData[2]["headers"])
    #         print('\n\n---------------------------------------------------------------------')
    #         break




